import { type TreeNode } from "./lib/types";
import { cn } from "@/src/utils/tailwind";
import {
  type ScoreDomain,
  ObservationLevel,
  type ObservationLevelType,
} from "@langfuse/shared";
import { useMemo, useRef, useLayoutEffect, useCallback } from "react";
import { InfoIcon, ChevronRight } from "lucide-react";
import { Button } from "@/src/components/ui/button";
import { ItemBadge } from "@/src/components/ItemBadge";
import { usePostHogClientCapture } from "@/src/features/posthog-analytics/usePostHogClientCapture";
import {
  calculateDisplayTotalCost,
  unnestObservation,
} from "@/src/components/trace/lib/helpers";
import type Decimal from "decimal.js";
import { SpanItem } from "@/src/components/trace/SpanItem";
import { type WithStringifiedMetadata } from "@/src/utils/clientSideDomainTypes";
import { api } from "@/src/utils/api";
import { useVirtualizer } from "@tanstack/react-virtual";

type FlatNode = {
  node: TreeNode;
  depth: number;
  treeLines: boolean[];
  isLastSibling: boolean;
};

const flattenTree = (
  node: TreeNode,
  collapsedNodes: string[],
  depth = 0,
  treeLines: boolean[] = [],
  isLastSibling = true,
): FlatNode[] => {
  const flatList: FlatNode[] = [{ node, depth, treeLines, isLastSibling }];

  if (node.children.length > 0 && !collapsedNodes.includes(node.id)) {
    const sortedChildren = node.children.sort(
      (a, b) => a.startTime.getTime() - b.startTime.getTime(),
    );
    sortedChildren.forEach((child, index) => {
      const isChildLast = index === sortedChildren.length - 1;

      flatList.push(
        ...flattenTree(
          child,
          collapsedNodes,
          depth + 1,
          [...treeLines, !isChildLast],
          isChildLast,
        ),
      );
    });
  }
  return flatList;
};

export const TraceTree = ({
  tree,
  collapsedNodes,
  toggleCollapsedNode,
  displayScores: scores,
  currentNodeId,
  setCurrentNodeId,
  showDuration,
  showCostTokens,
  showScores,
  colorCodeMetrics,
  nodeCommentCounts,
  className,
  showComments,
  hiddenObservationsCount,
  minLevel,
  setMinLevel,
  projectId,
  traceId,
}: {
  tree: TreeNode;
  collapsedNodes: string[];
  toggleCollapsedNode: (id: string) => void;
  displayScores: WithStringifiedMetadata<ScoreDomain>[];
  currentNodeId: string | undefined;
  setCurrentNodeId: (id: string | undefined) => void;
  showDuration: boolean;
  showCostTokens: boolean;
  showScores: boolean;
  colorCodeMetrics: boolean;
  nodeCommentCounts?: Map<string, number>;
  className?: string;
  showComments: boolean;
  hiddenObservationsCount?: number;
  minLevel?: ObservationLevelType;
  setMinLevel?: React.Dispatch<React.SetStateAction<ObservationLevelType>>;
  projectId: string;
  traceId: string;
}) => {
  const utils = api.useUtils();
  const parentRef = useRef<HTMLDivElement>(null);

  const handleObservationHover = useCallback(
    (node: TreeNode) => {
      if (node.type !== "TRACE") {
        void utils.observations.byId.prefetch(
          {
            observationId: node.id,
            startTime: node.startTime,
            traceId,
            projectId,
          },
          {
            staleTime: 5 * 60 * 1000,
          },
        );
      }
    },
    [utils, traceId, projectId],
  );

  const totalCost = useMemo(() => {
    const convertTreeNodeToObservation = (node: TreeNode): any => ({
      ...node,
      children: node.children.map(convertTreeNodeToObservation),
    });

    if (tree.type === "TRACE") {
      const allObservations = tree.children.flatMap((child) =>
        unnestObservation(convertTreeNodeToObservation(child)),
      );
      return calculateDisplayTotalCost({ allObservations });
    }

    return calculateDisplayTotalCost({
      allObservations: [convertTreeNodeToObservation(tree)],
    });
  }, [tree]);

  const flattenedItems = useMemo(() => {
    // The root node has depth 0, empty treeLines, and isLastSibling=true (it's the only root)
    return flattenTree(tree, collapsedNodes, 0, [], true);
  }, [tree, collapsedNodes]);

  // Calculate dynamic row height based on node content
  const estimateSize = useCallback(
    (index: number) => {
      const item = flattenedItems[index];
      if (!item) return 37;

      // Base height for node row
      let height = 37;

      // Add height for metrics line if visible
      const hasMetrics =
        (showDuration && (item.node.endTime || item.node.latency !== null)) ||
        (showCostTokens &&
          (item.node.inputUsage ||
            item.node.outputUsage ||
            item.node.totalUsage ||
            item.node.totalCost));
      if (hasMetrics) {
        height += 16; // One line of metrics
      }

      // Add height for scores if visible
      if (showScores) {
        const nodeScores = scores.filter((s) =>
          item.node.type === "TRACE"
            ? s.observationId === null
            : s.observationId === item.node.id,
        );

        if (nodeScores.length > 0) {
          // Estimate 3 score groups per line, each line ~20px
          const estimatedLines = Math.min(Math.ceil(nodeScores.length / 3), 3);
          height += estimatedLines * 20;
        }
      }

      return height;
    },
    [flattenedItems, showDuration, showCostTokens, showScores, scores],
  );

  const rowVirtualizer = useVirtualizer({
    count: flattenedItems.length,
    getScrollElement: () => parentRef.current,
    estimateSize,
    overscan: 500,
    measureElement:
      typeof window !== "undefined"
        ? (element) => element.getBoundingClientRect().height
        : undefined,
  });

  // Track initial currentNodeId to detect URL-based navigation vs user clicks
  const initialCurrentNodeIdRef = useRef(currentNodeId);
  const hasScrolledOnInitialLoadRef = useRef(false);

  // Handle scrolling to selected observation on initial page load only
  useLayoutEffect(() => {
    // Only scroll if:
    // 1. We have a currentNodeId
    // 2. We haven't scrolled yet
    // 3. The currentNodeId matches the initial one (meaning it came from URL, not user click)
    if (
      currentNodeId &&
      !hasScrolledOnInitialLoadRef.current &&
      currentNodeId === initialCurrentNodeIdRef.current
    ) {
      // Find the index of the item with the matching node ID
      const index = flattenedItems.findIndex(
        (item) => item.node.id === currentNodeId,
      );

      if (index !== -1) {
        // Use virtualizer's scrollToIndex to scroll to the item
        // This works even if the item isn't rendered yet
        rowVirtualizer.scrollToIndex(index, {
          align: "center",
          behavior: "smooth",
        });
        hasScrolledOnInitialLoadRef.current = true;
      }
    }
  }, [currentNodeId, flattenedItems, rowVirtualizer]);

  return (
    <div
      ref={parentRef}
      className={cn("h-full overflow-y-auto", className)}
      id="trace-tree"
    >
      <div
        style={{
          height: `${rowVirtualizer.getTotalSize()}px`,
          width: "100%",
          position: "relative",
        }}
      >
        {rowVirtualizer.getVirtualItems().map((virtualRow) => {
          const item = flattenedItems[virtualRow.index];
          return (
            <div
              key={item.node.id}
              data-index={virtualRow.index}
              ref={rowVirtualizer.measureElement}
              style={{
                position: "absolute",
                top: 0,
                left: 0,
                width: "100%",
                transform: `translateY(${virtualRow.start}px)`,
              }}
            >
              <TraceTreeRow
                node={item.node}
                depth={item.depth}
                treeLines={item.treeLines}
                isLastSibling={item.isLastSibling}
                collapsedNodes={collapsedNodes}
                toggleCollapsedNode={toggleCollapsedNode}
                scores={scores}
                comments={nodeCommentCounts}
                currentNodeId={currentNodeId}
                setCurrentNodeId={setCurrentNodeId}
                showDuration={showDuration}
                showCostTokens={showCostTokens}
                showScores={showScores}
                colorCodeMetrics={colorCodeMetrics}
                parentTotalCost={totalCost}
                parentTotalDuration={
                  tree.latency ? tree.latency * 1000 : undefined
                }
                showComments={showComments}
                onObservationHover={handleObservationHover}
              />
            </div>
          );
        })}
      </div>

      {minLevel && hiddenObservationsCount && hiddenObservationsCount > 0 ? (
        <span className="flex items-center gap-1 p-2 py-4">
          <InfoIcon className="h-4 w-4 text-muted-foreground" />
          <p className="text-sm text-muted-foreground">
            <span>
              {hiddenObservationsCount}{" "}
              {hiddenObservationsCount === 1 ? "observation" : "observations"}{" "}
              below {minLevel} level are hidden.{" "}
            </span>
            <span
              className="cursor-pointer underline"
              onClick={() => setMinLevel?.(ObservationLevel.DEBUG)}
            >
              Show all
            </span>
          </p>
        </span>
      ) : null}
    </div>
  );
};

type TraceTreeRowProps = {
  node: TreeNode;
  depth: number;
  treeLines: boolean[];
  isLastSibling: boolean;
  collapsedNodes: string[];
  toggleCollapsedNode: (id: string) => void;
  scores: WithStringifiedMetadata<ScoreDomain>[];
  comments?: Map<string, number>;
  currentNodeId: string | undefined;
  setCurrentNodeId: (id: string | undefined) => void;
  showDuration: boolean;
  showCostTokens: boolean;
  showScores: boolean;
  colorCodeMetrics: boolean;
  parentTotalCost?: Decimal;
  parentTotalDuration?: number;
  showComments: boolean;
  onObservationHover: (node: TreeNode) => void;
};

const TraceTreeRow = ({
  node,
  depth,
  treeLines,
  isLastSibling,
  collapsedNodes,
  toggleCollapsedNode,
  scores,
  comments,
  currentNodeId,
  setCurrentNodeId,
  showDuration,
  showCostTokens,
  showScores,
  colorCodeMetrics,
  parentTotalCost,
  parentTotalDuration,
  showComments,
  onObservationHover,
}: TraceTreeRowProps) => {
  const capture = usePostHogClientCapture();
  const collapsed = collapsedNodes.includes(node.id);

  const isSelected =
    currentNodeId === node.id || (!currentNodeId && node.type === "TRACE");

  return (
    <div
      className={cn(
        "relative flex w-full cursor-pointer rounded-md px-0",
        isSelected ? "bg-muted" : "hover:bg-muted/50",
      )}
      style={{
        paddingTop: 0,
        paddingBottom: 0,
        borderRadius: "0.5rem",
      }}
      onClick={(e) => {
        if (!e.currentTarget?.closest("[data-expand-button]")) {
          setCurrentNodeId(node.type === "TRACE" ? undefined : node.id);
        }
      }}
    >
      <div className="flex w-full pl-2">
        {/* 1. Indents: ancestor level indicators */}
        {depth > 0 && (
          <div className="flex flex-shrink-0">
            {Array.from({ length: depth - 1 }, (_, i) => (
              <div key={i} className="relative w-5">
                {treeLines[i] && (
                  <div className="absolute bottom-0 left-3 top-0 w-px bg-border" />
                )}
              </div>
            ))}
          </div>
        )}

        {/* 2. Current element bars: up/down/horizontal connectors */}
        {depth > 0 && (
          <div className="relative w-5 flex-shrink-0">
            <>
              {/* Vertical bar connecting upwards */}
              <div
                className={cn(
                  "absolute left-3 top-0 w-px bg-border",
                  isLastSibling ? "h-3" : "bottom-3",
                )}
              />
              {/* Vertical bar connecting downwards if not last sibling */}
              {!isLastSibling && (
                <div className="absolute bottom-0 left-3 top-3 w-px bg-border" />
              )}
              {/* Horizontal bar connecting to icon */}
              <div className="absolute left-3 top-3 h-px w-2 bg-border" />
            </>
          </div>
        )}

        {/* 3. Icon + child connector: fixed width container */}
        <div className="relative flex w-6 flex-shrink-0 flex-col py-1.5">
          <div className="relative z-10 flex h-4 items-center justify-center">
            <ItemBadge type={node.type} isSmall className="!size-3" />
          </div>
          {/* Vertical bar downwards if there are expanded children */}
          {node.children.length > 0 && !collapsed && (
            <div className="absolute bottom-0 left-1/2 top-3 w-px bg-border" />
          )}
          {/* Root node downward connector */}
          {depth === 0 && node.children.length > 0 && !collapsed && (
            <div className="absolute bottom-0 left-1/2 top-3 w-px bg-border" />
          )}
        </div>

        {/* 4. Content button: just the text/metrics content */}
        <button
          type="button"
          onClick={(e) => {
            e.stopPropagation();
            setCurrentNodeId(node.type === "TRACE" ? undefined : node.id);
          }}
          onMouseEnter={() => onObservationHover(node)}
          title={node.name}
          className={cn(
            "peer relative flex min-w-0 flex-1 items-start rounded-md py-0.5 pl-1 pr-2 text-left",
          )}
        >
          <SpanItem
            node={node}
            scores={scores}
            comments={comments}
            showDuration={showDuration}
            showCostTokens={showCostTokens}
            showScores={showScores}
            colorCodeMetrics={colorCodeMetrics}
            parentTotalCost={parentTotalCost}
            parentTotalDuration={parentTotalDuration}
            showComments={showComments}
          />
        </button>

        {/* 5. Expand/Collapse button */}
        {node.children.length > 0 && (
          <div className="flex items-center justify-end py-1 pr-1">
            <Button
              data-expand-button
              size="icon"
              variant="ghost"
              onClick={(ev) => {
                ev.stopPropagation();
                toggleCollapsedNode(node.id);
                capture(
                  collapsed
                    ? "trace_detail:observation_tree_expand"
                    : "trace_detail:observation_tree_collapse",
                  { type: "single", nodeType: node.type },
                );
              }}
              className="h-6 w-6 flex-shrink-0 hover:bg-primary/10"
            >
              <span
                className={cn(
                  "inline-block h-4 w-4 transform transition-transform duration-200 ease-in-out",
                  collapsed ? "rotate-0" : "rotate-90",
                )}
              >
                <ChevronRight className="h-4 w-4" />
              </span>
            </Button>
          </div>
        )}
      </div>
    </div>
  );
};
